#ifndef __MATMUL_H
#define __MATMUL_H

#include "operation.h"

namespace mslite {

template <typename Dtype>
class Matmul: public Operation<Dtype> {
public:
    int tiling_;
    Matmul(std::string name, std::vector<int> lhsShape, std::vector<int> rhsShape, int tiling);     // the shape order is HW
    int bind(void) override;
    int setInputs(const std::vector<std::vector<Dtype>>& inputs) override;
    std::vector<std::vector<Dtype>> getOutput(void) override;
};

template <typename Dtype>
Matmul<Dtype>::Matmul(std::string name, std::vector<int> lhsShape, std::vector<int> rhsShape, int tiling): Operation<Dtype>(name), tiling_(tiling) {
    this->global_.resize(2);
    this->local_.resize(2);
    this->global_ = {static_cast<size_t>(lhsShape[0]), static_cast<size_t>(rhsShape[1])};
    this->local_ = {1, 1};

    MemType otype = (name.find("ob") == std::string::npos) ? Image2d : Buffer;
    MemType ltype = (name.find("lb") == std::string::npos) ? Image2d : Buffer;
    MemType rtype = (name.find("rb") == std::string::npos) ? Image2d : Buffer;
    this->outputs_.emplace_back("out", std::vector<int>{lhsShape[0], rhsShape[1]}, HW, otype);
    std::swap(rhsShape[0], rhsShape[1]);
    this->inputs_.emplace_back("lhs", lhsShape, HW, ltype);
    this->inputs_.emplace_back("rhs", rhsShape, WH, rtype);
    this->bind();
}

template <typename Dtype>
int Matmul<Dtype>::bind(void) {
    cl_int status;
    std::vector<int> outputShape = this->outputs_[0].shape();
    std::vector<int> lhsShape = this->inputs_[0].shape();
    cl_int4 hwkt[] = {outputShape[0], outputShape[1], lhsShape[1], tiling_};
    status = clSetKernelArg(this->kernel_, 0, sizeof(cl_mem), static_cast<void *>(&this->outputs_[0].memObj()));
    CHECK(status);
    status = clSetKernelArg(this->kernel_, 1, sizeof(cl_mem), static_cast<void *>(&this->inputs_[0].memObj()));
    status = clSetKernelArg(this->kernel_, 2, sizeof(cl_mem), static_cast<void *>(&this->inputs_[1].memObj()));
    status = clSetKernelArg(this->kernel_, 3, sizeof(cl_int4), static_cast<void *>(&hwkt));
    return status;
}

template<typename Dtype>
int Matmul<Dtype>::setInputs(const std::vector<std::vector<Dtype>>& inputs) {
    cl_int status;
    assert(inputs.size() == this->inputs_.size());
    status = this->inputs_[0].setData(inputs[0], HW, tiling_); // lhs
    status = this->inputs_[1].setData(inputs[1], HW, tiling_); // rhs
    return status;
}

template<typename Dtype>
std::vector<std::vector<Dtype>> Matmul<Dtype>::getOutput(void) {
    std::vector<std::vector<Dtype>> res;
    for (auto output: this->outputs_) {
        res.push_back(output.getData());
    }
    return std::move(res);
}
    
}

#endif